package com.zyw.datasource;

import org.springframework.beans.BeansException;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;

import javax.sql.DataSource;
import java.util.Map;

/**
 * @author zhangyw
 * @date 2017/3/8 9:40
 */
public class DynamicDataSourceRegister implements ApplicationContextAware, InitializingBean {

    private ApplicationContext applicationContext;

    private DynamicDataSource dynamicDataSource;

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    public void registerDataSource(Object lookupKey, DataSource dataSource) {
        assert this.dynamicDataSource != null;

        this.dynamicDataSource.registerDataSource(lookupKey, dataSource);
    }

    public void registerDataSource(Map<Object, DataSource> dataSourceMap) {
        assert this.dynamicDataSource != null;
        assert dataSourceMap != null;
        for (Map.Entry<Object, DataSource> entry : dataSourceMap.entrySet()) {
            this.dynamicDataSource.registerDataSource(entry.getKey(), entry.getValue());
        }
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        this.dynamicDataSource = applicationContext.getBean(DynamicDataSource.class);

        if (this.dynamicDataSource == null){
            throw new IllegalStateException("dynamic datasource hasn't been initial");
        }
    }
}
